/*
 * COPYRIGHT (C) 2005-2011
 * RE2, INC.
 * ALL RIGHTS RESERVED
 *
 *
 * THIS WORK CONTAINS VALUABLE CONFIDENTIAL AND PROPRIETARY INFORMATION.
 * DISCLOSURE OR REPRODUCTION WITHOUT THE WRITTEN AUTHORIZATION OF RE2, INC.
 * IS PROHIBITED. THIS UNPUBLISHED WORK BY RE2, INC. IS PROTECTED BY THE LAWS
 * OF THE UNITED STATES AND OTHER COUNTRIES. IF PUBLICATION OF THE WORK SHOULD
 * OCCUR, THE FOLLOWING NOTICE SHALL APPLY.
 *
 * "COPYRIGHT (C) 2005-2011 RE2, INC. ALL RIGHTS RESERVED."
 *
 * RE2, INC. DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING
 * ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS, IN NO EVENT SHALL
 * RE2, INC. BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR
 * ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER
 * IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT
 * OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 */


#include <simple_sphere_detector/PointCloudProcessing.h>
#include <ros/ros.h>
#include "pcl/filters/voxel_grid.h"
#include "pcl/common/common_headers.h"
#include <iostream>


using namespace std;
using namespace pcl;


/**
 * Constructor takes in a node handle to allow subscribe/publish
 */
PointCloudProcessing::PointCloudProcessing(ros::NodeHandle node) :
		m_ptCloudTopic("/SR4000/points2"),
		m_running(false) {
    init(node);
}

/**
 * Initialize variables and the control ui
 */
void PointCloudProcessing::init(ros::NodeHandle node) {

    showFilteredPoints = true;
    findCylinderPoints = false;
    findSpherePoints = true;
    findPlanePoints = false;

    m_node = node;

    m_originalPointsPub = m_node.advertise<sensor_msgs::PointCloud2> ("OriginalPoints", 0);
    m_filteredPointsPub = m_node.advertise<sensor_msgs::PointCloud2> ("FilteredPoints", 0);
    m_inlierPointsPub = m_node.advertise<sensor_msgs::PointCloud2> ("InlierPoints", 0);
    m_sphereCoefficientsPub = m_node.advertise<geometry_msgs::QuaternionStamped> ("SphereCoefficients", 0);

    viewDepthMin = -0.01;
    viewDepthMax = 2.5;
    viewWidthMin = -1.25;
    viewWidthMax = 1.25;
    viewHeightMin = -0.8;
    viewHeightMax = 0.75;

    viewDepthMin  = -0.5;  // z-axis in SR4000 coordinates
    viewDepthMax  =  0.5;
    viewWidthMin  =  0.8;  // x-axis in SR4000 coordinates
    viewWidthMax  =  1.5;
    viewHeightMin = -0.5;  // y-axis in SR4000  coordinates
    viewHeightMax =  0.5;

    m_node.getParam("pointCloudTopic", m_ptCloudTopic);

    // Subscribe to the points
    m_kinectPointCloud2Sub = m_node.subscribe<sensor_msgs::PointCloud2>(m_ptCloudTopic, 1, &PointCloudProcessing::receivePointCloud2, this);
}

/**
 * Main thread needed to enable user ctrl-c shutdown
 */
void PointCloudProcessing::start() {

	m_running = true;

	m_mainThread = new boost::thread(&PointCloudProcessing::mainLoop,this);
}

/**
 *
 */
void PointCloudProcessing::mainLoop() {
	while (m_running) {
		ros::Duration(0.01).sleep();
		ros::spinOnce();
	}
}

/**
 *
 */
void PointCloudProcessing::stop() {

	m_originalPointsPub.shutdown();
	m_filteredPointsPub.shutdown();
	m_inlierPointsPub.shutdown();
	m_sphereCoefficientsPub.shutdown();
	m_kinectPointCloud2Sub.shutdown();

	if (m_running) {
		m_running = false;
	}
}


/**
 * Point cloud callback
 */
void PointCloudProcessing::receivePointCloud2(const sensor_msgs::PointCloud2::ConstPtr &msg) {

    pcl::fromROSMsg(*msg, m_origCloud);
    filterPoints(m_origCloud);

	if(m_cloudFiltered.points.size() > 0){
		if(findPlanePoints){
			findPlane(m_cloudFiltered);
		}

		if(findCylinderPoints){
			findCylinder(m_cloudFiltered);
		}

		if(findSpherePoints){
			findSphere(m_cloudFiltered);
		}
    } 

    publishPoints();
}

/**
 * Filter point clouds by depth. Take the closest "sphere".
 */
void PointCloudProcessing::filterPoints(pcl::PointCloud<PointType> cloud){
    sensor_msgs::PointCloud2 temp;
    pcl::PointCloud<PointType> cloud_temp;
    pcl::PointCloud<PointType> cloud_temp2;
    pcl::PointCloud<PointType> cloud_temp3;

    // Filter by distance
    pcl::PassThrough<PointType> pass;

    // Build a passthrough filter to remove points we don't want
    pass.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    pass.setFilterFieldName("z");
    pass.setFilterLimits(viewDepthMin, viewDepthMax); // Meters?   3.4 good for a while, then changes modes?
    pass.filter(cloud_temp3);

    pcl::PassThrough<PointType> pass2;
    pass2.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud_temp3));
    pass2.setFilterFieldName("y");
    pass2.setFilterLimits(viewHeightMin, viewHeightMax); // Meters? Flip (min is up, max is down)
    pass2.filter(cloud_temp2);

    pcl::PassThrough<PointType> pass3;
    pass3.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud_temp2));
    pass3.setFilterFieldName("x");
    pass3.setFilterLimits(viewWidthMin, viewWidthMax); // Meters?
    pass3.filter(cloud_temp);

    m_cloudFiltered = cloud_temp;
}

/**
 * Invert all points on X-axis in point cloud
 */
void PointCloudProcessing::flipX(pcl::PointCloud<PointType> *cloud){
    for (size_t i = 0; i < cloud->points.size(); ++i){
        cloud->points[i].x = cloud->points[i].x * -1.0;
    }
}

/**
 * Invert all points on Y-axis in point cloud
 */
void PointCloudProcessing::flipY(pcl::PointCloud<PointType> *cloud){
    for (size_t i = 0; i < cloud->points.size(); ++i){
        cloud->points[i].y = cloud->points[i].y * -1.0;
    }
}

/**
 * Invert all points on Z-axis in point cloud
 */
void PointCloudProcessing::flipZ(pcl::PointCloud<PointType> *cloud){
    for (size_t i = 0; i < cloud->points.size(); ++i){
        cloud->points[i].z = cloud->points[i].z * -1.0;
    }
}

/**
 * Find a plane or flat surface
 */
void PointCloudProcessing::findPlane(pcl::PointCloud<PointType> cloud){
    pcl::ModelCoefficients coefficients;
    pcl::PointIndices inliers;
    pcl::ExtractIndices<PointType> extract;

    // Create the segmentation object
    pcl::SACSegmentation<PointType> seg;

    // Optional
    seg.setOptimizeCoefficients(true);

    // Mandatory
    seg.setModelType(pcl::SACMODEL_PLANE);
    seg.setMethodType(pcl::SAC_RANSAC);
    seg.setDistanceThreshold(0.01);

    seg.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    seg.segment(inliers, coefficients);

    if (inliers.indices.size() == 0) {
        ROS_ERROR ("Could not estimate a planar model for the given dataset.");
    }

    m_cloudInliers.points.clear();

    // Extract the planar inliers from the input cloud
    extract.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    extract.setIndices(boost::make_shared<pcl::PointIndices>(inliers));
    extract.setNegative(false);
    extract.filter(m_cloudInliers);
}

/**
 * Find the closest cylinder in the point cloud
 */
void PointCloudProcessing::findCylinder(pcl::PointCloud<PointType> cloud){

    // All the objects needed
    pcl::NormalEstimation<PointType, pcl::Normal> ne;
    pcl::SACSegmentationFromNormals<PointType, pcl::Normal> seg;
    pcl::ExtractIndices<PointType> extract;
    pcl::ExtractIndices<pcl::Normal> extract_normals;
    pcl::KdTreeANN<PointType>::Ptr tree = boost::make_shared<pcl::KdTreeANN<PointType> >();

    // Datasets
    pcl::PointCloud<pcl::Normal> cloud_normals;
    pcl::ModelCoefficients coefficients_plane, coefficients_cylinder;
    pcl::PointIndices inliers_plane, inliers_cylinder;


    // Estimate point normals
    ne.setSearchMethod(tree);
    ne.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    ne.setKSearch(50);
    ne.compute(cloud_normals);

    // Create the segmentation object for the planar model and set all the parameters
    seg.setOptimizeCoefficients(true);
    seg.setModelType(pcl::SACMODEL_NORMAL_PLANE);
    seg.setMethodType(pcl::SAC_RANSAC);
    seg.setNormalDistanceWeight(0.1);
    seg.setMaxIterations(100);
    seg.setDistanceThreshold(0.03);
    seg.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    seg.setInputNormals(boost::make_shared<pcl::PointCloud<pcl::Normal> >(cloud_normals));

    // Obtain the plane inliers and coefficients
    seg.segment(inliers_plane, coefficients_plane);

    // Extract the planar inliers from the input cloud
    extract.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    extract.setIndices(boost::make_shared<pcl::PointIndices>(inliers_plane));
    extract.setNegative(false);

    // Write the planar inliers to disk
    pcl::PointCloud<PointType> cloud_plane;
    extract.filter(cloud_plane);

    // Remove the planar inliers, extract the rest
    extract.setNegative(true);
    extract.filter(cloud);
    extract_normals.setNegative(true);
    extract_normals.setInputCloud(boost::make_shared<pcl::PointCloud<pcl::Normal> >(cloud_normals));
    extract_normals.setIndices(boost::make_shared<pcl::PointIndices>(inliers_plane));
    extract_normals.filter(cloud_normals);

    // Create the segmentation object for cylinder segmentation and set all the parameters
    seg.setOptimizeCoefficients(true);
    seg.setModelType(pcl::SACMODEL_CYLINDER);
    seg.setMethodType(pcl::SAC_RANSAC);
    seg.setNormalDistanceWeight(0.1);
    seg.setMaxIterations(10000);
    seg.setDistanceThreshold(0.02);
    seg.setRadiusLimits(0.0025, 0.025);
    seg.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    seg.setInputNormals(boost::make_shared<pcl::PointCloud<pcl::Normal> >(cloud_normals));

    // Obtain the cylinder inliers and coefficients
    seg.segment(inliers_cylinder, coefficients_cylinder);

    // Write the cylinder inliers to disk
    extract.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    extract.setIndices(boost::make_shared<pcl::PointIndices>(inliers_cylinder));
    extract.setNegative(false);
    m_cloudInliers.points.clear();
    extract.filter(m_cloudInliers);
}

/**
 * Find the closest sphere in the point cloud
 */
void PointCloudProcessing::findSphere(pcl::PointCloud<PointType> cloud){

    // All the objects needed
    pcl::NormalEstimation<PointType, pcl::Normal> ne;
    pcl::SACSegmentationFromNormals<PointType, pcl::Normal> seg;
    pcl::ExtractIndices<PointType> extract;
    pcl::ExtractIndices<pcl::Normal> extract_normals;
    pcl::KdTreeANN<PointType>::Ptr tree = boost::make_shared<pcl::KdTreeANN<PointType> >();

    // Datasets
    pcl::PointCloud<pcl::Normal> cloud_normals;
    pcl::ModelCoefficients coefficients_plane, coefficients_cylinder;
    pcl::PointIndices inliers_plane, inliers_cylinder;

    // Estimate point normals
    ne.setSearchMethod(tree);
    ne.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    //ne.setKSearch(50);
    ne.setKSearch(std::min(cloud.size(), (size_t)50));
    ne.compute(cloud_normals);

    // Create the segmentation object for the planar model and set all the parameters
    seg.setOptimizeCoefficients(true);
    seg.setModelType(pcl::SACMODEL_NORMAL_PLANE);
    seg.setMethodType(pcl::SAC_RANSAC);
    seg.setNormalDistanceWeight(0.1);
    seg.setMaxIterations(100);
    seg.setDistanceThreshold(0.03);
    seg.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    seg.setInputNormals(boost::make_shared<pcl::PointCloud<pcl::Normal> >(cloud_normals));

    // Obtain the plane inliers and coefficients
    seg.segment(inliers_plane, coefficients_plane);

    // Extract the planar inliers from the input cloud
    extract.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    extract.setIndices(boost::make_shared<pcl::PointIndices>(inliers_plane));
    extract.setNegative(false);

    // Write the planar inliers to disk
    pcl::PointCloud<PointType> cloud_plane;
    extract.filter(cloud_plane);

    // Remove the planar inliers, extract the rest
    extract.setNegative(true);
    extract.filter(cloud);
    extract_normals.setNegative(true);
    extract_normals.setInputCloud(boost::make_shared<pcl::PointCloud<pcl::Normal> >(cloud_normals));
    extract_normals.setIndices(boost::make_shared<pcl::PointIndices>(inliers_plane));
    extract_normals.filter(cloud_normals);

    PointCloud<PointType> cluster_points;

    // Obtain the Sphere inliers and coefficients
    ModelCoefficients coefficients_sphere;
    PointIndices inliers_sphere;
    SACSegmentation<PointType> seg_sphere;
    seg_sphere.setOptimizeCoefficients(true);
    seg_sphere.setMethodType(SAC_RANSAC);
    seg_sphere.setModelType(SACMODEL_SPHERE);
    seg_sphere.setRadiusLimits(0.01, 0.05);
    seg_sphere.setMaxIterations(500);
    seg_sphere.setDistanceThreshold(0.005);
    seg_sphere.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));

    // Process goes too slow when there are more points
    if (cloud.size() > 1800 || cloud.size() < 50) return;

    seg_sphere.segment(inliers_sphere, coefficients_sphere);

    ROS_DEBUG_STREAM("Sphere coefficients: " << coefficients_sphere);

    // Write the cylinder inliers to disk
    extract.setInputCloud(boost::make_shared<pcl::PointCloud<PointType> >(cloud));
    extract.setIndices(boost::make_shared<pcl::PointIndices>(inliers_sphere));
    extract.setNegative(false);
    m_cloudInliers.points.clear();
    extract.filter(m_cloudInliers);

    {
      geometry_msgs::QuaternionStamped msg;

      msg.header.frame_id = coefficients_sphere.header.frame_id;
      msg.quaternion.x    = coefficients_sphere.values[0];
      msg.quaternion.y    = coefficients_sphere.values[1];
      msg.quaternion.z    = coefficients_sphere.values[2];
      msg.quaternion.w    = coefficients_sphere.values[3];  // radius

      m_sphereCoefficientsPub.publish(msg);
    }
}

/**
 *
 */
void PointCloudProcessing::publishPoints(){

    // Original points
    if(m_origCloud.size() > 0){
        sensor_msgs::PointCloud2 msg;
        pcl::toROSMsg(m_origCloud,msg);

        m_originalPointsPub.publish(msg);
    }

    // Filtered points
    if(m_cloudFiltered.size() > 0){
        sensor_msgs::PointCloud2 msg;
        pcl::toROSMsg(m_cloudFiltered,msg);

        m_filteredPointsPub.publish(msg);
    }

    // Object points
    if(m_cloudInliers.size() > 0){
        sensor_msgs::PointCloud2 msg;
        pcl::toROSMsg(m_cloudInliers,msg);

        m_inlierPointsPub.publish(msg);
    }

}


/**
 * Main
 */
int main(int argc, char *argv[]) {
    ros::init(argc, argv, "mimic");
    ros::NodeHandle node("");

    boost::shared_ptr<PointCloudProcessing> pcp(new PointCloudProcessing(node));

    pcp->start();

    ros::waitForShutdown();

    pcp->stop();
}
